import os
import sys
from collections import Counter

import fire
import pandas as pd
from zkl_aiutils_datasets import load_dataset
from zkl_pyutils_fsspec import FsLike, resolve_fs

# Set up project path
project_dir_path = os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../"))
sys.path.append(project_dir_path)

# TODO: 重构为包中的一部分

def main(
    fs: FsLike = os.path.join(project_dir_path, "datasets/ready/ukbiobank/v4.1"),
    output_csv: str = os.path.join(project_dir_path, "counts/ukbiobank/v4.1/str_counts.csv"),
    type_name: str = "str_columns",
):
    """
    Counts the occurrences of each unique string in the specified dataset column
    and saves the results to a CSV file.
    """
    print("Starting string count process...")

    # Load the dataset
    fs = resolve_fs(fs)
    dataset = load_dataset(fs)
    type_dataset = dataset.named_children[type_name]
    print(f"Dataset loaded from: {fs}")

    # Count strings
    string_counts = Counter()
    print("Counting strings...")
    try:
        num_samples = len(type_dataset)
    except TypeError:
        num_samples = "unknown"

    for i, sample in enumerate(type_dataset):
        if (i + 1) % 10000 == 0:
            if num_samples != "unknown":
                print(f"Processing sample {i + 1}/{num_samples}")
            else:
                print(f"Processing sample {i + 1}")

        for text in sample:
            if text is not None:
                string_counts[text] += 1
    
    print("Finished counting strings.")

    # Prepare data for CSV output
    print("Preparing data for CSV output...")
    # Convert counter to a list of tuples (already sorted by most common)
    count_list = string_counts.most_common()
    
    # Create a DataFrame
    df = pd.DataFrame(count_list, columns=['string', 'count'])

    # Ensure output directory exists
    output_dir = os.path.dirname(output_csv)
    os.makedirs(output_dir, exist_ok=True)
    print(f"Output directory '{output_dir}' ensured.")

    # Save to CSV
    df.to_csv(output_csv, index=False)
    print(f"Successfully saved results to {output_csv}")
    print(f"Total unique strings found: {len(df)}")


if __name__ == '__main__':
    fire.Fire(main)